import os
import numpy as np
import torch
import matplotlib.pyplot as plt

from model import build_resnet_18_model
from preprocess import preprocess_audio, extract_features  # 保持不变

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
DEVICE = torch.device('cpu')
TARGET_SAMPLE_RATE = 22050
MODEL_PATH = '../models/best_model.pth'
LABELS = {0: '非病理性', 1: '哮喘', 2: '咽炎', 3: '支气管炎', 4: '百日咳'}

N_MELS = 128  # Mel谱图特征数量
N_MFCC = 13   # MFCC特征数量

def load_trained_model(model_path, input_shape):
    model = build_resnet_18_model(input_shape)
    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
    model.to(DEVICE)
    model.eval()
    return model

def predict_cough_type(model, feature):
    # feature 原始形状假定为 (H, W) 或 (H, W, channels)
    # 如果是二维特征，则添加通道和 batch 维度
    if feature.ndim == 2:
        feature = feature[np.newaxis, np.newaxis, :, :]
    elif feature.ndim == 3:
        # 如果最后一个维度为 1，则转为 (1, 1, H, W)
        if feature.shape[-1] == 1:
            feature = np.transpose(feature, (2, 0, 1))[np.newaxis, ...]
        else:
            # 若通道数不为 1，根据实际情况调整
            feature = np.transpose(feature, (2, 0, 1))[np.newaxis, ...]
    tensor_feature = torch.tensor(feature, dtype=torch.float32).to(DEVICE)
    with torch.no_grad():
        outputs = model(tensor_feature)
        _, predicted = torch.max(outputs, 1)
    return LABELS[predicted.item()]

def plot_features(feature):
    plt.figure(figsize=(14, 6))
    # Mel 谱图特征
    mel_feature = feature[:, :N_MELS]
    plt.subplot(1, 2, 1)
    plt.imshow(mel_feature.T, aspect='auto', origin='lower')
    plt.title('Mel Spectrogram Feature')
    plt.xlabel('Time')
    plt.ylabel('Frequency')
    plt.colorbar(format='%+2.0f dB')
    # MFCC 特征
    mfcc_feature = feature[:, N_MELS:N_MELS + N_MFCC]
    plt.subplot(1, 2, 2)
    plt.imshow(mfcc_feature.T, aspect='auto', origin='lower')
    plt.title('MFCC Feature')
    plt.xlabel('Time')
    plt.ylabel('MFCC Coefficients')
    plt.colorbar()
    plt.tight_layout()
    plt.savefig("temp/extracted_features.png")
    plt.show()

def main(wav_file_path):
    # 预处理音频文件
    preprocessed_audio = preprocess_audio(wav_file_path)
    # 提取特征
    feature = extract_features(preprocessed_audio, TARGET_SAMPLE_RATE)
    plot_features(feature)
    # 假定 feature 的形状为 (H, W, channels) 或 (H, W)
    # 确定输入形状（转换为 (channels, H, W)）
    if feature.ndim == 2:
        input_shape = (1, feature.shape[0], feature.shape[1])
    elif feature.ndim == 3:
        if feature.shape[-1] == 1:
            input_shape = (1, feature.shape[0], feature.shape[1])
        else:
            # 若有多个通道，根据实际情况设置
            input_shape = (feature.shape[-1], feature.shape[0], feature.shape[1])
    model = load_trained_model(MODEL_PATH, input_shape)
    cough_type = predict_cough_type(model, feature)
    print(f'The cough type is: {cough_type}')

if __name__ == '__main__':
    wav_file_path = r"D:\python\cough_recognition\data\raw\百日咳\百日咳_14.wav"
    main(wav_file_path)
